from flask import Flask, jsonify, send_from_directory, request
import pandas as pd
import numpy as np
from scipy import stats
import os
import joblib
from data_processor import DataProcessor
from database import db  # 导入数据库实例
from flask_cors import CORS  # 添加CORS支持

# 获取当前文件的目录
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# 创建Flask应用，指定静态文件目录
app = Flask(__name__, static_folder=os.path.join(BASE_DIR, 'visualizations'))
app.static_url_path = '/static'  # 设置静态文件URL前缀
app.static_folder = os.path.join(BASE_DIR, 'visualizations')  # 设置静态文件目录
CORS(app)  # 启用CORS

# 添加JS文件路由
@app.route('/js/<path:filename>')
def serve_js(filename):
    return send_from_directory(os.path.join(BASE_DIR, 'visualizations', 'js'), filename)

# 初始化数据库连接
print("初始化数据库连接...")
try:
    db.connect()
    print(f"数据库已连接: {db.db_path}")
except Exception as e:
    print(f"数据库连接错误: {str(e)}")

# 加载数据
def load_data():
    data = pd.read_csv(os.path.join(BASE_DIR, 'data', 'heart_disease_full.csv'))
    return data

# 加载模型
def load_model():
    model = joblib.load(os.path.join(BASE_DIR, 'models', 'best_model.joblib'))
    return model

# 初始化数据处理器
data_processor = DataProcessor({})

# 关闭数据库连接
@app.teardown_appcontext
def shutdown_session(exception=None):
    db.close()

# 服务静态文件
@app.route('/')
def index():
    return send_from_directory(os.path.join(BASE_DIR, 'visualizations'), 'index.html')

@app.route('/predict')
def predict_page():
    return send_from_directory(os.path.join(BASE_DIR, 'visualizations'), 'predict.html')

@app.route('/api/class_distribution')
def get_class_distribution():
    data = load_data()
    class_counts = data['target'].value_counts().sort_index()

    option = {
        'title': {
            'text': '心脏病患者分布',
            'left': 'center'
        },
        'tooltip': {
            'trigger': 'item',
            'formatter': '{b}: {c} ({d}%)'
        },
        'legend': {
            'orient': 'vertical',
            'left': 'left',
            'data': ['无心脏病', '有心脏病']
        },
        'series': [{
            'name': '患者分布',
            'type': 'pie',
            'radius': '50%',
            'data': [
                {'value': int(class_counts[0]), 'name': '无心脏病'},
                {'value': int(class_counts[1]), 'name': '有心脏病'}
            ],
            'emphasis': {
                'itemStyle': {
                    'shadowBlur': 10,
                    'shadowOffsetX': 0,
                    'shadowColor': 'rgba(0, 0, 0, 0.5)'
                }
            }
        }]
    }

    return jsonify(option)

@app.route('/api/correlation_heatmap')
def get_correlation_heatmap():
    data = load_data()
    corr = data.corr()

    option = {
        'title': {
            'text': '特征相关性热力图',
            'left': 'center'
        },
        'tooltip': {
            'position': 'top'
        },
        'grid': {
            'height': '60%',
            'top': '10%'
        },
        'xAxis': {
            'type': 'category',
            'data': list(corr.columns),
            'splitArea': {
                'show': True
            },
            'axisLabel': {
                'interval': 0,
                'rotate': 45
            }
        },
        'yAxis': {
            'type': 'category',
            'data': list(corr.columns),
            'splitArea': {
                'show': True
            }
        },
        'visualMap': {
            'min': -1,
            'max': 1,
            'calculable': True,
            'orient': 'horizontal',
            'left': 'center',
            'bottom': '15%'
        },
        'series': [{
            'name': '相关性',
            'type': 'heatmap',
            'data': [[i, j, round(corr.iloc[i, j], 2)] for i in range(len(corr.columns)) for j in range(len(corr.columns))],
            'label': {
                'show': True
            },
            'emphasis': {
                'itemStyle': {
                    'shadowBlur': 10,
                    'shadowColor': 'rgba(0, 0, 0, 0.5)'
                }
            }
        }]
    }

    return jsonify(option)

@app.route('/api/target_correlation')
def get_target_correlation():
    data = load_data()
    corr = data.corr()['target'].sort_values(ascending=True)
    corr = corr[corr.index != 'target']

    option = {
        'title': {
            'text': '特征与心脏病的相关性',
            'left': 'center'
        },
        'tooltip': {
            'trigger': 'axis',
            'axisPointer': {
                'type': 'shadow'
            }
        },
        'grid': {
            'left': '3%',
            'right': '4%',
            'bottom': '15%',
            'containLabel': True
        },
        'xAxis': {
            'type': 'value',
            'name': '相关系数'
        },
        'yAxis': {
            'type': 'category',
            'data': list(corr.index),
            'axisLabel': {
                'interval': 0
            }
        },
        'series': [{
            'name': '相关系数',
            'type': 'bar',
            'data': [round(x, 3) for x in corr.values],
            'label': {
                'show': True,
                'position': 'right'
            }
        }]
    }

    return jsonify(option)

@app.route('/api/feature_distribution/<feature>')
def get_feature_distribution(feature):
    data = load_data()

    if feature not in data.columns:
        return jsonify({'error': 'Feature not found'}), 404

    feature_data = data[feature]

    if feature in ['sex', 'cp']:
        # 分类特征
        counts = feature_data.value_counts().sort_index()
        labels = {
            'sex': ['女性', '男性'],
            'cp': ['典型心绞痛', '非典型心绞痛', '非心绞痛', '无症状']
        }

        option = {
            'title': {
                'text': f'{feature} 分布',
                'left': 'center'
            },
            'tooltip': {
                'trigger': 'item'
            },
            'legend': {
                'orient': 'vertical',
                'left': 'left'
            },
            'series': [{
                'name': feature,
                'type': 'pie',
                'radius': '50%',
                'data': [
                    {'value': int(v), 'name': labels[feature][i]}
                    for i, v in enumerate(counts)
                ]
            }]
        }
    else:
        # 数值特征
        hist, bins = np.histogram(feature_data, bins='auto')
        bin_centers = (bins[:-1] + bins[1:]) / 2

        option = {
            'title': {
                'text': f'{feature} 分布',
                'left': 'center'
            },
            'tooltip': {
                'trigger': 'axis'
            },
            'xAxis': {
                'type': 'category',
                'data': [round(x, 2) for x in bin_centers]
            },
            'yAxis': {
                'type': 'value',
                'name': '频数'
            },
            'series': [{
                'name': feature,
                'type': 'bar',
                'data': hist.tolist()
            }]
        }

    return jsonify(option)

@app.route('/api/predict', methods=['POST'])
def predict():
    try:
        # 获取请求数据
        if request.is_json:
            data = request.get_json()
            print("接收到的JSON数据:", data)  # 调试日志
        else:
            # 处理表单数据
            data = request.form.to_dict()
            print("接收到的表单数据:", data)  # 调试日志

            # 转换数据类型
            for key in data:
                if key == 'user_id':
                    continue  # 保留user_id为字符串
                elif key in ['age', 'trestbps', 'chol', 'thalach', 'oldpeak']:
                    try:
                        data[key] = float(data[key])
                    except ValueError:
                        return jsonify({
                            'error': f'无效的{key}值',
                            'message': f'{key}必须是数字'
                        }), 400
                else:
                    try:
                        data[key] = int(data[key])
                    except ValueError:
                        return jsonify({
                            'error': f'无效的{key}值',
                            'message': f'{key}必须是整数'
                        }), 400

        if not data:
            print("错误：没有接收到数据")  # 调试日志
            return jsonify({
                'error': '没有接收到数据',
                'message': '请提供预测所需的数据'
            }), 400

        # 获取用户ID (如果有)
        user_id = data.pop('user_id', None)

        # 验证必需字段
        required_fields = ['age', 'sex', 'cp', 'trestbps', 'chol', 'fbs', 'restecg',
                         'thalach', 'exang', 'oldpeak', 'slope', 'ca', 'thal']
        missing_fields = [field for field in required_fields if field not in data]
        if missing_fields:
            print(f"错误：缺少必需字段: {missing_fields}")  # 调试日志
            return jsonify({
                'error': f'缺少必需字段: {", ".join(missing_fields)}',
                'message': '请填写所有必需的字段'
            }), 400

        # 转换为DataFrame
        input_data = pd.DataFrame([data])
        print("转换后的DataFrame:", input_data)  # 调试日志

        # 确保输入数据经过离群值去除和缺失值填补
        input_data = data_processor.remove_outliers(input_data)
        input_data = data_processor.fill_missing_values(input_data)

        # 数据预处理
        try:
            processed_data = data_processor.prepare_prediction_data(input_data)
            print("预处理后的数据:", processed_data)  # 调试日志
        except Exception as e:
            print("数据预处理错误:", str(e))  # 调试日志
            return jsonify({
                'error': str(e),
                'message': '数据预处理失败，请检查输入数据格式'
            }), 400

        # 加载模型
        try:
            model_path = os.path.join(BASE_DIR, 'models', 'best_model.joblib')
            print(f"尝试加载模型: {model_path}")  # 调试日志
            if not os.path.exists(model_path):
                print(f"错误：模型文件不存在: {model_path}")  # 调试日志
                return jsonify({
                    'error': '模型文件不存在',
                    'message': '请先训练模型'
                }), 500
            model = joblib.load(model_path)
            print("模型加载成功")  # 调试日志
        except Exception as e:
            print("模型加载错误:", str(e))  # 调试日志
            return jsonify({
                'error': str(e),
                'message': '模型加载失败，请检查模型文件'
            }), 500

        # 进行预测
        try:
            prediction = model.predict(processed_data)[0]
            probability = model.predict_proba(processed_data)[0][1]
            print(f"预测结果: {prediction}, 概率: {probability}")  # 调试日志

            # 保存预测结果到数据库
            # 确保数据库连接已初始化
            if db.connection is None:
                db.connect()

            prediction_id = db.save_prediction(int(prediction), float(probability), data, user_id)
            print(f"预测结果已保存到数据库，ID: {prediction_id}，用户ID: {user_id}")

        except Exception as e:
            print("预测错误:", str(e))  # 调试日志
            return jsonify({
                'error': str(e),
                'message': '预测失败，请检查输入数据'
            }), 400

        return jsonify({
            'prediction': int(prediction),
            'probability': float(probability),
            'prediction_id': prediction_id,
            'user_id': user_id
        })

    except Exception as e:
        print("系统错误:", str(e))  # 调试日志
        return jsonify({
            'error': str(e),
            'message': '系统错误，请稍后重试'
        }), 500

# 获取历史预测记录
@app.route('/api/prediction_history', methods=['GET'])
def get_prediction_history():
    try:
        limit = request.args.get('limit', default=10, type=int)
        # 确保数据库连接已初始化
        if db.connection is None:
            db.connect()

        records = db.get_recent_predictions(limit)

        # 转换记录为JSON格式
        result = []
        for record in records:
            result.append({
                'id': record[0],
                'prediction': record[1],
                'probability': record[2],
                'timestamp': record[3]
            })

        return jsonify({
            'success': True,
            'records': result
        })
    except Exception as e:
        print("获取历史预测记录错误:", str(e))
        return jsonify({
            'error': str(e),
            'message': '获取历史预测记录失败'
        }), 500

# 通过用户ID查询预测记录
@app.route('/api/predictions/user/<user_id>', methods=['GET'])
def get_predictions_by_user_id(user_id):
    try:
        if not user_id:
            return jsonify({
                'error': '用户ID不能为空',
                'message': '请提供有效的用户ID'
            }), 400

        # 确保数据库连接已初始化
        if db.connection is None:
            db.connect()

        records = db.get_prediction_by_user_id(user_id)

        if not records:
            return jsonify({
                'success': True,
                'message': f'未找到用户ID为{user_id}的预测记录',
                'records': []
            })

        # 转换记录为JSON格式
        result = []
        for i, record in enumerate(records):
            result.append({
                'id': record[0],
                'user_id': record[1],
                'prediction': record[2],
                'probability': record[3],
                'timestamp': record[4],
                'sequence': i + 1,  # 添加序号
                'input_data': {
                    'age': record[5],
                    'sex': record[6],
                    'cp': record[7],
                    'trestbps': record[8],
                    'chol': record[9],
                    'fbs': record[10],
                    'restecg': record[11],
                    'thalach': record[12],
                    'exang': record[13],
                    'oldpeak': record[14],
                    'slope': record[15],
                    'ca': record[16],
                    'thal': record[17]
                }
            })

        return jsonify({
            'success': True,
            'total_records': len(result),
            'user_id': user_id,
            'records': result
        })
    except Exception as e:
        print("获取用户预测记录错误:", str(e))
        return jsonify({
            'error': str(e),
            'message': '获取用户预测记录失败'
        }), 500

if __name__ == '__main__':
    app.run(debug=True) 